package io.openkunlun.javarpc.codegen;

import com.google.common.collect.Lists;
import freemarker.template.Configuration;
import freemarker.template.Template;
import freemarker.template.TemplateExceptionHandler;
import io.openkunlun.javarpc.RpcCall;
import io.openkunlun.javarpc.RpcIgnore;
import io.openkunlun.javarpc.RpcService;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.ElementFilter;
import javax.tools.Diagnostic;
import java.io.Writer;
import java.util.*;
import java.util.stream.Collectors;

public class RpcGen extends AbstractProcessor {

    @Override
    public synchronized void init(ProcessingEnvironment processingEnvironment) {
        super.init(processingEnvironment);
    }

    @Override
    public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnvironment) {

        Collection<? extends Element> annotatedElements = roundEnvironment.getElementsAnnotatedWith(RpcService.class);
        List<TypeElement> types = ElementFilter.typesIn(annotatedElements);
        if (types.isEmpty()) {
            return false;
        }

        Configuration cfg = new Configuration(Configuration.VERSION_2_3_23);
        cfg.setClassForTemplateLoading(RpcGen.class, "/templates/");
        cfg.setDefaultEncoding("UTF-8");
        cfg.setLocale(Locale.CHINESE);
        cfg.setTemplateExceptionHandler(TemplateExceptionHandler.RETHROW_HANDLER);

        for (TypeElement element : types) {
            String className = element.getSimpleName().toString();
            if (element.getKind() == ElementKind.INTERFACE) {
                RpcService rpcService = element.getAnnotation(RpcService.class);
                if (rpcService.generateHandler()) {
                    genHandler(element, className, cfg, rpcService);
                }
                if (rpcService.generateStub()) {
                    genStub(element, className, cfg, rpcService);
                }
            }
        }
        return true;
    }

    private void genHandler(TypeElement element, String className, Configuration cfg, RpcService rpcService) {
        String qualifiedName = element.getQualifiedName().toString();
        String service = rpcService.service();
        String contentType = rpcService.contentType();

        RpcClass rpcClass = RpcClass.newBuilder()
                .app(rpcService.app())
                .service(service.isEmpty() ? qualifiedName : service)
                .contentType(contentType)
                .packageName(qualifiedName.substring(0, qualifiedName.lastIndexOf(".")))
                .className(className)
                .build();
        List<RpcMethod> rpcMethods = newRpcMethods(element);

        try (Writer writer = processingEnv.getFiler().createSourceFile(qualifiedName + "Handler").openWriter()) {
            Map<String, Object> dataModel = new HashMap<>();
            dataModel.put("rpcClass", rpcClass);
            dataModel.put("rpcMethods", rpcMethods);

            Template template = cfg.getTemplate("RpcHandler.ftl");
            template.process(dataModel, writer);
        } catch (Exception e) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, e.getMessage());
        }
    }

    private void genStub(TypeElement element, String className, Configuration cfg, RpcService rpcService) {
        String qualifiedName = element.getQualifiedName().toString();
        String service = rpcService.service();
        String contentType = rpcService.contentType();

        RpcClass rpcClass = RpcClass.newBuilder()
                .app(rpcService.app())
                .service(service.isEmpty() ? qualifiedName : service)
                .contentType(contentType)
                .packageName(qualifiedName.substring(0, qualifiedName.lastIndexOf(".")))
                .className(className)
                .build();
        List<RpcMethod> rpcMethods = newRpcMethods(element);

        try (Writer writer = processingEnv.getFiler().createSourceFile(qualifiedName + "Stub").openWriter()) {
            Map<String, Object> dataMap = new HashMap<>();
            dataMap.put("rpcClass", rpcClass);
            dataMap.put("rpcMethods", rpcMethods);

            Template template = cfg.getTemplate("RpcStub.ftl");
            template.process(dataMap, writer);
        } catch (Exception e) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, e.getMessage());
        }
    }

    private List<RpcMethod> newRpcMethods(TypeElement element) {
        List<RpcMethod> results = Lists.newArrayList();
        for (Element k : element.getEnclosedElements()) {
            ExecutableElement executableElement = (ExecutableElement) k;
            RpcCall rpcCall = k.getAnnotation(RpcCall.class);
            RpcIgnore rpcIgnore = k.getAnnotation(RpcIgnore.class);
            List<RpcParameter> parameters = newRpcParameters(executableElement);
            List<String> exceptions = newRpcExceptions(executableElement);
            RpcMethod rpcMethod = RpcMethod.newBuilder()
                    .ignored(null != rpcIgnore)
                    .methodId(null != rpcCall ? rpcCall.method() : k.getSimpleName().toString())
                    .methodName(k.getSimpleName().toString())
                    .returnType(executableElement.getReturnType().toString())
                    .returnCastType(RpcUtils.castType(executableElement.getReturnType().toString()))
                    .parameters(parameters)
                    .exceptions(exceptions)
                    .build();
            results.add(rpcMethod);
        }

        return results;
    }

    private List<RpcParameter> newRpcParameters(ExecutableElement element) {
        return element
                .getParameters()
                .stream()
                .map(this::newRpcParameter)
                .collect(Collectors.toList());
    }

    private RpcParameter newRpcParameter(VariableElement it) {
        return RpcParameter
                .newBuilder()
                .type(it.asType().toString())
                .castType(RpcUtils.castType(it.asType().toString()))
                .name(it.getSimpleName().toString())
                .build();
    }

    private List<String> newRpcExceptions(ExecutableElement element){
        return element.getThrownTypes().stream().map(TypeMirror::toString).collect(Collectors.toList());
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Collections.singleton(RpcService.class.getCanonicalName());
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.RELEASE_8;
    }
}